import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.Token;
import org.antlr.v4.runtime.tree.ParseTreeWalker;

import java.util.*;

// 把字节码转换为机器码
public class 字节码汇编器 extends AssemblerBaseListener {

    public static 字节码汇编器 汇编(String 字节码) {
        // 1. 构造抽象语法树
        var input = CharStreams.fromString(字节码);
        var lexer = new AssemblerLexer(input);// 词法分析器，将输入文本转换为词法符号
        var tokens = new CommonTokenStream(lexer);// 词法符号流
        var parser = new AssemblerParser(tokens);// 语法分析器

        var 结果 = new 字节码汇编器(字节码定义.指令集);
        ParseTreeWalker.DEFAULT.walk(结果, parser.program());
        return 结果;
    }


    public static final int 初始化代码尺寸 = 1024;

    protected Map<String, Integer> 操作码指令映射 = new HashMap<>(); // 将操作码映射为指令ID

    protected Map<String, 标签符号> 标签 = new HashMap<>(); // 标签作用域

    protected List<Object> 常量池 = new ArrayList<>(); // 在常量池中，所有的浮点数和字符串都有唯一的索引

    protected int 指令指针 = 0; // 指令地址指针（IP）；用于填写代码

    protected byte[] 机器码 = new byte[初始化代码尺寸]; // 代码内存

    protected int 数据尺寸; // 通过 .global 设置

    protected 函数符号 入口函数; // main 函数

    public 字节码汇编器(字节码定义.指令[] 指令集) {
        for (int i = 1; i < 指令集.length; i++) { // 注意这里是从1开始，因为0表示无效地址
            操作码指令映射.put(指令集[i].名称.toLowerCase(), i);
        }
    }

    public byte[] 取机器码() {return 机器码;}

    public int 取机器码尺寸() {return 指令指针;}

    public int 取数据尺寸() {return 数据尺寸;}

    public 函数符号 取入口函数() {return 入口函数;}


    @Override
    public void enterInstr_noArgs(AssemblerParser.Instr_noArgsContext ctx) {
        生成指令(ctx.ID().getSymbol());
    }

    @Override
    public void enterInstr_oneArgs(AssemblerParser.Instr_oneArgsContext ctx) {
        生成指令(ctx.ID().getSymbol());
        生成操作数(ctx.operand().start);
    }

    @Override
    public void enterInstr_twoArgs(AssemblerParser.Instr_twoArgsContext ctx) {
        生成指令(ctx.ID().getSymbol());
        生成操作数(ctx.a.start);
        生成操作数(ctx.b.start);
    }

    @Override
    public void enterInstr_threeArgs(AssemblerParser.Instr_threeArgsContext ctx) {
        生成指令(ctx.ID().getSymbol());
        生成操作数(ctx.a.start);
        生成操作数(ctx.b.start);
        生成操作数(ctx.c.start);
    }

    protected void 生成指令(Token 指令Token) {
        var 名称 = 指令Token.getText();
        var 指令码 = 操作码指令映射.get(名称);
        if (指令码 == null) {
            System.err.println("[第" + 指令Token.getLine() + "行]: 未知指令 " + 名称);
            return;
        }

        自动扩容(指令指针 + 1);
        机器码[指令指针] = (byte) (指令码 & 0xFF); // 把指令码限制在 [0,255] 范围内
        指令指针++;
    }

    protected void 生成操作数(Token 操作数Token) {
        var 文本 = 操作数Token.getText();
        int 值 = 字节码定义.无效;
        switch (操作数Token.getType()) {
            case AssemblerParser.INT -> 值 = Integer.parseInt(文本);
            case AssemblerParser.CHAR -> 值 = 文本.charAt(1); // 去掉单引号
            case AssemblerParser.FLOAT -> 值 = 取常量索引(Float.parseFloat(文本));
            case AssemblerParser.STRING -> 值 = 取常量索引(去引号(文本));
            case AssemblerParser.ID -> 值 = 取标签地址(文本);
            case AssemblerParser.FUNC -> 值 = 取函数索引(文本.substring(0, 文本.length() - 2));
            case AssemblerParser.REG -> 值 = 取寄存器(操作数Token);
        }

        自动扩容(指令指针 + 4);
        写整数(机器码, 指令指针, 值);
        指令指针 += 4;
    }

    protected int 取常量索引(Object o) {
        if (常量池.contains(o)) return 常量池.indexOf(o);
        常量池.add(o);
        return 常量池.size() - 1;
    }

    protected int 取标签地址(String 名称) {
        var 符号 = 标签.get(名称);
        if (符号 == null) {
            符号 = new 标签符号(名称, 指令指针, true);
            符号.是否已定义 = false;
            标签.put(名称, 符号);
        } else {
            if (符号.是否为前向引用) {
                符号.添加前向引用(指令指针);
            } else {
                return 符号.地址;
            }
        }

        return 0;
    }
    
    protected String 去引号(String 文本) {
        return 文本.substring(1, 文本.length() - 1);
    }

    @Override
    public void enterFunctionDeclaration(AssemblerParser.FunctionDeclarationContext ctx) {
        var 名称 = ctx.name.getText();
        var 参数个数 = Integer.parseInt(ctx.a.getText());
        var 局部变量个数 = Integer.parseInt(ctx.lo.getText());

        var 函数符号 = new 函数符号(名称, 参数个数, 局部变量个数, 指令指针);
        if (名称.equals("main")) {
            入口函数 = 函数符号;
        }

        if (常量池.contains(函数符号)) {
            常量池.set(常量池.indexOf(函数符号), 函数符号);
        } else {
            取常量索引(函数符号);
        }

    }

    protected int 取函数索引(String 名称) {
        var 索引 = 常量池.indexOf(new 函数符号(名称));
        if (索引 >= 0) {
            return 索引;
        }

        return 取常量索引(new 函数符号(名称));
    }

    private int 取寄存器(Token token) {
        // 转换 "rN" -> N
        var rs = token.getText();
        rs = rs.substring(1);
        return Integer.parseInt(rs);
    }

    private void 自动扩容(int 索引) {
        if (索引 >= 机器码.length) {
            var 新机器码 = new byte[机器码.length * 2];
            System.arraycopy(机器码, 0, 新机器码, 0, 机器码.length);
            机器码 = 新机器码;
        }
    }

    @Override
    public void enterGlobals(AssemblerParser.GlobalsContext ctx) {
        数据尺寸 = Integer.parseInt(ctx.INT().getText());
    }

    @Override
    public void enterLabel(AssemblerParser.LabelContext ctx) {
        var 名称 = ctx.ID().getText();
        var 符号 = 标签.get(名称);
        if (符号 == null) {
            符号 = new 标签符号(名称, 指令指针, false);
            标签.put(名称, 符号); // 没有则定义一个
        } else {
            if (符号.是否为前向引用) {
                符号.是否已定义 = true;
                符号.地址 = 指令指针;
                符号.解析前向引用(机器码);
            } else {
                System.err.println("[第" + ctx.ID().getSymbol().getLine() + "行]标签 " + 名称 + " 重复定义");
            }
        }
    }

    public static int 取整数(byte[] 内存, int 索引) {
        int b1 = 内存[索引++] & 0xFF; // mask off sign-extended bits
        int b2 = 内存[索引++] & 0xFF;
        int b3 = 内存[索引++] & 0xFF;
        int b4 = 内存[索引] & 0xFF;
        return b1 << (8 * 3) | b2 << (8 * 2) | b3 << (8) | b4;
    }

    public static void 写整数(byte[] 机器码, int 索引, int 值) {
        机器码[索引] = (byte) ((值 >> (8 * 3)) & 0xFF); // get highest byte  // 大端模式
        机器码[索引 + 1] = (byte) ((值 >> (8 * 2)) & 0xFF);
        机器码[索引 + 2] = (byte) ((值 >> (8)) & 0xFF);
        机器码[索引 + 3] = (byte) (值 & 0xFF);
    }
}
